/* tag::catalog[]
Title:: Payload Builder Size Tests

Goal:: Test the consensus payload builder and the accompaning payload validator.

Runbook::
. Set up two subnets with one fast node each
. Install a universal canister in both, one is called target canister the other assist canister.
. The assist canister will be used to send the xnet data to the target canister.
. Send ingress message to target canister, that is slightly below maximum size. Expect it to succeed.

Success:: The payload builder respects the boundaries set by the registry, while the payload validator
accepts all payloads generated by the payload builder.

Coverage::
. The maximum size of an individual ingress message is respected.

end::catalog[] */

use anyhow::Result;
use assert_matches::assert_matches;
use ic_agent::AgentError;
use ic_agent::agent_error::HttpErrorPayload;
use ic_consensus_system_test_utils::rw_message::install_nns_and_check_progress;
use ic_registry_subnet_type::SubnetType;
use ic_system_test_driver::{
    driver::{
        group::SystemTestGroup,
        ic::InternetComputer,
        test_env::TestEnv,
        test_env_api::{HasPublicApiUrl, HasTopologySnapshot, IcNodeContainer, IcNodeSnapshot},
    },
    systest,
    util::UniversalCanister,
};
use ic_universal_canister::wasm;

const MAX_QUERY_MESSAGE_SIZE_BYTES: u64 = 4 * 1024 * 1024;
const REQUEST_HEADERS_OVERHEAD: u64 = 360;
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Status/413
const CONTENT_TOO_LARGE_STATUS: u16 = 413;

/// Sets up a testnet with
/// 1. System subnet with a single node
/// 2. Application subnet with a single node
/// 3. Boundary node
///
/// and installs [`UniversalCanister`]s on the subnets.
fn setup(env: TestEnv) {
    InternetComputer::new()
        .add_fast_single_node_subnet(SubnetType::System)
        .add_fast_single_node_subnet(SubnetType::Application)
        .with_api_boundary_nodes(1)
        .setup_and_start(&env)
        .expect("failed to setup IC under test");

    install_nns_and_check_progress(env.topology_snapshot());

    for api_bn in env.topology_snapshot().api_boundary_nodes() {
        api_bn
            .await_status_is_healthy()
            .expect("API boundary node did not come up healthy.");
    }
}

fn get_first_node(env: &TestEnv, subnet_type: SubnetType) -> IcNodeSnapshot {
    env.topology_snapshot()
        .subnets()
        .find(|subnet| subnet.subnet_type() == subnet_type)
        .expect("There should be at least one subnet for every subnet type")
        .nodes()
        .next()
        .expect("Every subnet should have at least one node")
}

enum RequestType {
    Query,
    Update,
}

/// Makes an ingress call to the specified canister with a message of the
/// specified size.
async fn make_ingress_call(
    canister: &UniversalCanister<'_>,
    size: usize,
    request: RequestType,
) -> Result<Vec<u8>, AgentError> {
    match request {
        RequestType::Query => {
            canister
                .query(wasm().reply().stable_write(0, &vec![0; size]))
                .await
        }
        RequestType::Update => {
            canister
                .update(wasm().reply().stable_write(0, &vec![0; size]))
                .await
        }
    }
}

fn send_request(
    env: &TestEnv,
    subnet_type: SubnetType,
    message_size: u64,
    call: RequestType,
) -> Result<(), AgentError> {
    let node = get_first_node(env, subnet_type);

    // get an agent for the API boundary node
    let api_bn = env
        .topology_snapshot()
        .api_boundary_nodes()
        .next()
        .expect("There should be at least one API boundary node");
    let agent = api_bn.build_default_agent();

    let logger = env.logger();

    tokio::runtime::Runtime::new()
        .unwrap()
        .block_on(async move {
            let canister = UniversalCanister::new_with_params_with_retries(
                &agent,
                node.effective_canister_id(),
                /*compute_allocation= */ None,
                /*cycles= */ None,
                /*pages= */ Some(100),
                &logger,
            )
            .await;

            make_ingress_call(&canister, message_size as usize, call)
                .await
                .map(|_| ())
        })
}

fn test_app_subnet_update_within_limits(env: TestEnv) {
    assert_eq!(
        send_request(
            &env,
            SubnetType::Application,
            ic_limits::MAX_INGRESS_BYTES_PER_MESSAGE_APP_SUBNET - REQUEST_HEADERS_OVERHEAD,
            RequestType::Update,
        ),
        Ok(())
    );
}

fn test_app_subnet_query_within_limits(env: TestEnv) {
    assert_eq!(
        send_request(
            &env,
            SubnetType::Application,
            MAX_QUERY_MESSAGE_SIZE_BYTES - REQUEST_HEADERS_OVERHEAD,
            RequestType::Query,
        ),
        Ok(())
    );
}

fn test_app_subnet_update_exceeds_limits(env: TestEnv) {
    assert_matches!(
        send_request(
            &env,
            SubnetType::Application,
            ic_limits::MAX_INGRESS_BYTES_PER_MESSAGE_APP_SUBNET + 1,
            RequestType::Update,
        ),
        Err(AgentError::HttpError(HttpErrorPayload {
            status: CONTENT_TOO_LARGE_STATUS,
            ..
        }))
    );
}

fn test_app_subnet_query_exceeds_limits(env: TestEnv) {
    assert_matches!(
        send_request(
            &env,
            SubnetType::Application,
            MAX_QUERY_MESSAGE_SIZE_BYTES + 1,
            RequestType::Query,
        ),
        Err(AgentError::HttpError(HttpErrorPayload {
            status: CONTENT_TOO_LARGE_STATUS,
            ..
        }))
    );
}

fn test_nns_subnet_update_within_limits(env: TestEnv) {
    assert_eq!(
        send_request(
            &env,
            SubnetType::System,
            ic_limits::MAX_INGRESS_BYTES_PER_MESSAGE_NNS_SUBNET - REQUEST_HEADERS_OVERHEAD,
            RequestType::Update,
        ),
        Ok(())
    );
}

fn test_nns_subnet_query_within_limits(env: TestEnv) {
    assert_eq!(
        send_request(
            &env,
            SubnetType::System,
            MAX_QUERY_MESSAGE_SIZE_BYTES - REQUEST_HEADERS_OVERHEAD,
            RequestType::Query,
        ),
        Ok(())
    );
}

fn test_nns_subnet_update_exceeds_limits(env: TestEnv) {
    assert_matches!(
        send_request(
            &env,
            SubnetType::System,
            ic_limits::MAX_INGRESS_BYTES_PER_MESSAGE_NNS_SUBNET + 1,
            RequestType::Update,
        ),
        Err(AgentError::HttpError(HttpErrorPayload {
            status: CONTENT_TOO_LARGE_STATUS,
            ..
        }))
    );
}

fn test_nns_subnet_query_exceeds_limits(env: TestEnv) {
    assert_matches!(
        send_request(
            &env,
            SubnetType::System,
            MAX_QUERY_MESSAGE_SIZE_BYTES + 1,
            RequestType::Query,
        ),
        Err(AgentError::HttpError(HttpErrorPayload {
            status: CONTENT_TOO_LARGE_STATUS,
            ..
        }))
    );
}

fn main() -> Result<()> {
    SystemTestGroup::new()
        .with_setup(setup)
        .add_test(systest!(test_app_subnet_update_within_limits))
        .add_test(systest!(test_app_subnet_update_exceeds_limits))
        .add_test(systest!(test_app_subnet_query_within_limits))
        .add_test(systest!(test_app_subnet_query_exceeds_limits))
        .add_test(systest!(test_nns_subnet_query_within_limits))
        .add_test(systest!(test_nns_subnet_query_exceeds_limits))
        .add_test(systest!(test_nns_subnet_update_within_limits))
        .add_test(systest!(test_nns_subnet_update_exceeds_limits))
        .execute_from_args()?;
    Ok(())
}
